package org.vison.wonfu.framework.soa.context;

import com.alibaba.ttl.TransmittableThreadLocal;

public class RpcContext {
    private static final TransmittableThreadLocal<ServiceContext> SERVICE_CONTEXT = new TransmittableThreadLocal<ServiceContext>(true) {
        @Override
        protected ServiceContext initialValue() {
            return new ServiceContext();
        }

        @Override
        protected ServiceContext childValue(ServiceContext parentValue) {
            return copy(parentValue);
        }

        @Override
        public ServiceContext copy(ServiceContext parentValue) {
            if (parentValue != null) {
                return parentValue.copy();
            }
            return super.copy(null);
        }
    };

    private static final TransmittableThreadLocal<ClientContextAttachment> CLIENT_ATTACHMENT = new TransmittableThreadLocal<ClientContextAttachment>(true) {
        @Override
        protected ClientContextAttachment initialValue() {
            return new ClientContextAttachment();
        }

        @Override
        protected ClientContextAttachment childValue(ClientContextAttachment parentValue) {
            return copy(parentValue);
        }

        @Override
        public ClientContextAttachment copy(ClientContextAttachment parentValue) {
            if (parentValue != null) {
                return parentValue.copy();
            }
            return super.copy(null);
        }
    };

    private static final TransmittableThreadLocal<ContextAttachment> SERVER_ATTACHMENT = new TransmittableThreadLocal<ContextAttachment>(true) {
        @Override
        protected ContextAttachment initialValue() {
            return new ContextAttachment();
        }

        @Override
        protected ContextAttachment childValue(ContextAttachment parentValue) {
            return copy(parentValue);
        }

        @Override
        public ContextAttachment copy(ContextAttachment parentValue) {
            if (parentValue != null) {
                return parentValue.copy();
            }
            return super.copy(null);
        }
    };

    public static ServiceContext getServiceContext() {
        return SERVICE_CONTEXT.get();
    }

    /**
     * 框架私有方法，业务请勿使用
     */
    public static void setServiceContext(ServiceContext serviceContext) {
        SERVICE_CONTEXT.set(serviceContext);
    }

    /**
     * get consumer side attachment ( A --> B , in A side)
     */
    public static ClientContextAttachment getClientAttachment() {
        return CLIENT_ATTACHMENT.get();
    }

    /**
     * get provider side attachment from consumer ( A --> B , in B side)
     */
    public static ContextAttachment getServerAttachment() {
        return SERVER_ATTACHMENT.get();
    }

    /**
     * 框架私有方法，业务请勿使用
     */
    public static void setServerAttachment(ContextAttachment contextAttachment) {
        SERVER_ATTACHMENT.set(contextAttachment);
    }

    public static void removeServiceContext() {
        SERVICE_CONTEXT.remove();
    }

    public static void removeClientAttachment() {
        CLIENT_ATTACHMENT.remove();
    }

    public static void removeServerAttachment() {
        SERVER_ATTACHMENT.remove();
    }

    public static void removeContext() {
        removeClientAttachment();
        removeServerAttachment();
        removeServiceContext();
    }

    public static RestoreServiceContext storeServiceContext() {
        return storeServiceContext(true);
    }

    public static RestoreServiceContext storeServiceContext(boolean removeCurrent) {
        RestoreServiceContext restoreServiceContext = new RestoreServiceContext();

        if (removeCurrent) {
            removeServiceContext();
        }

        return restoreServiceContext;
    }

    public static void restoreServiceContext(RestoreServiceContext restoreServiceContext) {
        if (restoreServiceContext != null) {
            restoreServiceContext.restore();
        }
    }

    /**
     * Restore ServiceContext
     */
    public static class RestoreServiceContext {
        private final ServiceContext serviceContext;

        public RestoreServiceContext() {
            this.serviceContext = SERVICE_CONTEXT.get();
        }

        protected void restore() {
            if (serviceContext != null) {
                SERVICE_CONTEXT.set(serviceContext);
            } else {
                removeServiceContext();
            }
        }

        public ServiceContext getServiceContext() {
            return serviceContext;
        }
    }

}
